#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np

# 特征数量，需与main.cpp中的NUM_FEATURES保持一致
NUM_FEATURES = 2048
# 动量参数
MOMENTUM = 0.9
# 防止方差为0的小值
EPSILON = 1e-5

def gen_golden_data():
    # 生成随机输入数据
    local_mean = np.random.uniform(0, 1, [NUM_FEATURES]).astype(np.float32)
    local_var = np.random.uniform(0.1, 1, [NUM_FEATURES]).astype(np.float32)
    global_mean = np.random.uniform(0, 1, [NUM_FEATURES]).astype(np.float32)
    global_var = np.random.uniform(0.1, 1, [NUM_FEATURES]).astype(np.float32)
    
    # 计算黄金标准（用于验证）
    golden_mean = MOMENTUM * global_mean + (1 - MOMENTUM) * local_mean
    golden_var = MOMENTUM * global_var + (1 - MOMENTUM) * local_var + EPSILON
    
    # 创建目录并保存数据
    import os
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    local_mean.tofile("./input/input_local_mean.bin")
    local_var.tofile("./input/input_local_var.bin")
    global_mean.tofile("./input/input_global_mean.bin")
    global_var.tofile("./input/input_global_var.bin")
    golden_mean.tofile("./output/output_golden_mean.bin")
    golden_var.tofile("./output/output_golden_var.bin")


if __name__ == "__main__":
    gen_golden_data()


